#include "wintun_driver.h"
#include <memory>
#include <winsock2.h>
#include <SetupAPI.h>
#include <cfgmgr32.h>
#include <Windows.h>
#include <ws2ipdef.h>
#include <iphlpapi.h>
#include <WS2tcpip.h>
#include "base/xlog.h"
#include "windivert.h"

static WINTUN_CREATE_ADAPTER_FUNC *WintunCreateAdapter;
static WINTUN_CLOSE_ADAPTER_FUNC *WintunCloseAdapter;
static WINTUN_OPEN_ADAPTER_FUNC *WintunOpenAdapter;
static WINTUN_GET_ADAPTER_LUID_FUNC *WintunGetAdapterLUID;
static WINTUN_GET_RUNNING_DRIVER_VERSION_FUNC *WintunGetRunningDriverVersion;
static WINTUN_DELETE_DRIVER_FUNC *WintunDeleteDriver;
static WINTUN_SET_LOGGER_FUNC *WintunSetLogger;
static WINTUN_START_SESSION_FUNC *WintunStartSession;
static WINTUN_END_SESSION_FUNC *WintunEndSession;
static WINTUN_GET_READ_WAIT_EVENT_FUNC *WintunGetReadWaitEvent;
static WINTUN_RECEIVE_PACKET_FUNC *WintunReceivePacket;
static WINTUN_RELEASE_RECEIVE_PACKET_FUNC *WintunReleaseReceivePacket;
static WINTUN_ALLOCATE_SEND_PACKET_FUNC *WintunAllocateSendPacket;
static WINTUN_SEND_PACKET_FUNC *WintunSendPacket;

/**
 * Wintun adapter descriptor.
 */
typedef struct HSWDEVICE__ *HSWDEVICE;
typedef struct _WINTUN_ADAPTER
{
    HSWDEVICE SwDevice;
    HDEVINFO DevInfo;
    SP_DEVINFO_DATA DevInfoData;
    WCHAR *InterfaceFilename;
    GUID CfgInstanceID;
    WCHAR DevInstanceID[MAX_DEVICE_ID_LEN];
    DWORD LuidIndex;
    DWORD IfType;
    DWORD IfIndex;
} WINTUN_ADAPTER;

static HMODULE InitializeWintun()
{
    HMODULE Wintun =
        LoadLibraryExW(L"wintun.dll", NULL, LOAD_LIBRARY_SEARCH_APPLICATION_DIR | LOAD_LIBRARY_SEARCH_SYSTEM32);
    if (!Wintun)
        return NULL;
#define X(Name) ((*(FARPROC *)&Name = GetProcAddress(Wintun, #Name)) == NULL)
    if (X(WintunCreateAdapter) || X(WintunCloseAdapter) || X(WintunOpenAdapter) || X(WintunGetAdapterLUID) ||
        X(WintunGetRunningDriverVersion) || X(WintunDeleteDriver) || X(WintunSetLogger) || X(WintunStartSession) ||
        X(WintunEndSession) || X(WintunGetReadWaitEvent) || X(WintunReceivePacket) || X(WintunReleaseReceivePacket) ||
        X(WintunAllocateSendPacket) || X(WintunSendPacket))
#undef X
    {
        DWORD LastError = GetLastError();
        FreeLibrary(Wintun);
        SetLastError(LastError);
        return NULL;
    }
    return Wintun;
}

WintunDriver::WintunDriver():
    queue_inject_(1024)
{

}

WintunDriver::~WintunDriver()
{
    is_close_.test_and_set();
    queue_inject_nonempty_.test_and_set();
    queue_inject_nonempty_.notify_all();
    SetEvent(event_close_);

    for (auto& thd : workers_) {
        thd.join();
    }

    CloseHandle(event_close_);
    WintunEndSession(session_);
    WintunCloseAdapter(adapter_);
    SINFO("WintunDriver end");
}

bool WintunDriver::init(cb_outbound_data func)
{
    cb_out_data_ = func;
    event_close_ = CreateEventW(NULL, TRUE, FALSE, NULL);
    if (!event_close_) {
        SERROR("Failed to create event_close", GetLastError());
        return false;
    }

    HMODULE Wintun = InitializeWintun();
    if (!Wintun) {
        SERROR("Failed to initialize Wintun,{}",GetLastError());
        return false;
    }
    SINFO("Wintun library loaded");

    GUID D2sGuid = { 0xdeadbabe, 0xcafe, 0xbeef, { 0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef } };
    adapter_ = WintunCreateAdapter(L"Drv2socks", L"Proxy", &D2sGuid);
    if (!adapter_) {
        SERROR("Failed to create adapter,{}", GetLastError());
        return false;
    }
    DWORD Version = WintunGetRunningDriverVersion();
    SINFO("Wintun v{}.{} loaded", (Version >> 16) & 0xff, (Version >> 0) & 0xff);

    MIB_UNICASTIPADDRESS_ROW AddressRow;
    InitializeUnicastIpAddressEntry(&AddressRow);
    WintunGetAdapterLUID(adapter_, &AddressRow.InterfaceLuid);
    AddressRow.Address.Ipv4.sin_family = AF_INET;
    AddressRow.Address.Ipv4.sin_addr.S_un.S_addr = htonl((10 << 24) | (6 << 16) | (7 << 8) | (7 << 0)); /* 10.6.7.7 */
    // AddressRow.Address.Ipv6.sin6_family = AF_INET6;
    // const uint8_t v6[16] = {0xfe,0x80,0xff,0xff};
    // WinDivertHelperHtonIpv6Address((const uint32_t*)&v6,(uint32_t*)&AddressRow.Address.Ipv6.sin6_addr.u);

    AddressRow.OnLinkPrefixLength = 32; /* This is a /24 network */
    AddressRow.DadState = IpDadStatePreferred;
    DWORD LastError = CreateUnicastIpAddressEntry(&AddressRow);
    if (LastError != ERROR_SUCCESS && LastError != ERROR_OBJECT_ALREADY_EXISTS) {
        WintunCloseAdapter(adapter_);
        SERROR("Failed to set IPv4 address,{}", LastError);
        return false;
    }
#if 1
    MIB_UNICASTIPADDRESS_ROW AddressRowV6;
    InitializeUnicastIpAddressEntry(&AddressRowV6);
    WintunGetAdapterLUID(adapter_, &AddressRowV6.InterfaceLuid);
    AddressRowV6.Address.Ipv6.sin6_family = AF_INET6;
    inet_pton(AF_INET6,"fe80:aabb::1",&AddressRowV6.Address.Ipv6.sin6_addr);
    AddressRowV6.OnLinkPrefixLength = 64; /* This is a /24 network */
    AddressRowV6.DadState = IpDadStatePreferred;
    LastError = CreateUnicastIpAddressEntry(&AddressRowV6);
    if (LastError != ERROR_SUCCESS && LastError != ERROR_OBJECT_ALREADY_EXISTS) {
        WintunCloseAdapter(adapter_);
        SERROR("Failed to set IPv6 address,{}", LastError);
        return false;
    }
    MIB_IPFORWARD_ROW2 forwardRow = { 0 };

    InitializeIpForwardEntry(&forwardRow);

    WintunGetAdapterLUID(adapter_, &forwardRow.InterfaceLuid);
    forwardRow.DestinationPrefix.Prefix.si_family = AF_INET6;
    forwardRow.DestinationPrefix.PrefixLength = 0; // Default route
    forwardRow.NextHop.si_family = AF_INET6;
    forwardRow.NextHop.Ipv6.sin6_family = AF_INET6;
    inet_pton(AF_INET6,"fe80::aaaa:4aff:fe2f:41d2",&forwardRow.NextHop.Ipv6.sin6_addr);
    forwardRow.SitePrefixLength = 0;
    forwardRow.Metric = 0; // Automatic metric

    DWORD result = CreateIpForwardEntry2(&forwardRow);
    if (result != ERROR_SUCCESS) {
        SERROR("Failed to set NextHop address,{}", LastError);
        // return false;
    }
#endif
    session_ = WintunStartSession(adapter_, 0x400000);
    if (!session_) {
        WintunCloseAdapter(adapter_);
        SERROR("Failed to create adapter");
        return false;
    }

    is_close_.clear();
    workers_.emplace_back(std::bind(&WintunDriver::sendPackets,this));
    workers_.emplace_back(std::bind(&WintunDriver::receivePackets,this));
    return true;
}

void WintunDriver::doWrite(NetPacket *buffer)
{
    if (is_close_.test()) return;

    queue_inject_.emplace(buffer);
    if (!queue_inject_nonempty_.test()) {
        queue_inject_nonempty_.test_and_set();
        queue_inject_nonempty_.notify_all();
    }
}

void WintunDriver::sendPackets()
{
    while (!is_close_.test()) {
        if (queue_inject_.empty()) {
            queue_inject_nonempty_.clear();
            queue_inject_nonempty_.wait(false,std::memory_order_relaxed);
            if (is_close_.test()) return;
        }

        NetPacket** it = queue_inject_.front();
        if (it == nullptr) {
            continue;
        }
        NetPacket* buffer = *it;
        queue_inject_.pop();

        BYTE *packet = WintunAllocateSendPacket(session_, buffer->data_len);
        if (packet) {
            std::memcpy(packet,buffer->data,buffer->data_len);
            WintunSendPacket(session_, packet);
            NetPacketPool::deletePacket(buffer);
        } else if (GetLastError() != ERROR_BUFFER_OVERFLOW) {
            SERROR("Packet write failed");
        }

    }
}

void WintunDriver::receivePackets()
{
    HANDLE wait_handles[] = { WintunGetReadWaitEvent(session_), event_close_ };
    while (!is_close_.test()) {
        DWORD packet_size;
        BYTE *Packet = WintunReceivePacket(session_, &packet_size);
        if (Packet) {
            NetPacket::Ptr netpk = _NetPacketPool->getSharedPacket(packet_size);
            std::memcpy(netpk->data,Packet,packet_size);
            cb_out_data_(netpk);
            WintunReleaseReceivePacket(session_, Packet);
        } else {
            DWORD last_error = GetLastError();
            switch (last_error)
            {
            case ERROR_NO_MORE_ITEMS:
                if (WaitForMultipleObjects(2, wait_handles, FALSE, INFINITE) == WAIT_OBJECT_0)
                    continue;
                return;
            default:
                SERROR("Packet read failed,{}",last_error);
                return;
            }
        }
    }
}
